from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_socketio import SocketIO, emit
import sqlite3
import json
from transformer import transformerVideoToAudio
import os
from ipex_llm.transformers import AutoModel
from transformers import AutoTokenizer
from langchain_community.document_loaders import DirectoryLoader
from langchain_community.llms import ChatGLM
from langchain.prompts import PromptTemplate
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
import websockets

db = None
qa = None
mc = None
sqa = None

# 加载嵌入模型
embedding_model_dict = {
    "ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
    "ernie-base": "nghuyong/ernie-3.0-base-zh",
    "text2vec": "GanymedeNil/text2vec-large-chinese",
    "text2vec2": "uer/sbert-base-chinese-nli",
    "text2vec3": "shibing624/text2vec-base-chinese",
}



def load_documents(directory = "contexts"):
    loader = DirectoryLoader(directory)
    documents = loader.load()
    text_spliter = CharacterTextSplitter(chunk_size=256, chunk_overlap=0)
    split_docs = text_spliter.split_documents(documents)
    return split_docs

def load_embedding_model(model_name="ernie-tiny"):
    encode_kwargs = {"normalize_embeddings": False}
    # 不需要为 CPU 使用指定设备
    return HuggingFaceEmbeddings(
        model_name=embedding_model_dict[model_name],
        encode_kwargs=encode_kwargs
    )

# 加载嵌入模型
embeddings = load_embedding_model('text2vec3')
# 加载或创建数据库

print("----------- Embedding模型加载成功 -----------")

print("-------------------------------------------------")


# 指定数据库文件的绝对路径
db_file_path = "D:/AtomGit/web_back/local/database/video_database.db"

# 连接到SQLite数据库（如果不存在，则会自动创建）
conn = sqlite3.connect(db_file_path)

# 创建游标对象
c = conn.cursor()

# 创建视频表
c.execute('''CREATE TABLE IF NOT EXISTS videos
             (id INTEGER PRIMARY KEY AUTOINCREMENT,
              name TEXT,
              tags TEXT,
              storage_path TEXT,
              nodejson TEXT,
              edgejson TEXT,
              treejson TEXT,
              content TEXT)''')

# 创建联查表
c.execute('''CREATE TABLE IF NOT EXISTS link
             (node TEXT,
              video_id INTEGER,
              FOREIGN KEY(video_id) REFERENCES videos(id))''')

# 提交更改并关闭连接
conn.commit()
conn.close()

# 创建语言模型
llm = ChatGLM(
    endpoint_url='http://127.0.0.1:8000',
    max_token=80000,
    top_p=0.9
)

def store_chroma(docs, embeddings, persist_directory):
    db = Chroma.from_documents(docs, embeddings, persist_directory=persist_directory)
    db.persist()
    return db


def check_and_create_vectorstore():
    # 初始文件夹名
    folder_name = "VectorStore"
    i = 1
    
    while True:
        # 当前文件夹路径
        current_folder = f"{folder_name}{i}"
        
        # 检查当前文件夹是否存在
        if not os.path.exists(current_folder):
            # 如果不存在，则创建该文件夹
            os.makedirs(current_folder)
            print(f"Created folder: {current_folder}")
            return current_folder  # 返回新创建的文件夹名
        
        # 如果当前文件夹已存在，则检查下一个文件夹
        i += 1

        
def text_get(text_path):
    global db
    global qa
    
    print("------------------{}--------------".format(text_path))
    
    documents = load_documents(text_path)
    # print(documents)
    db_dir = check_and_create_vectorstore()
    db = store_chroma(documents, embeddings, db_dir)
    
    # if not os.path.exists('VectorStore'):
    #     documents = load_documents(text_path)
    #     print(documents)
    #     db = store_chroma(documents, embeddings)
    # else:
    #     db = Chroma(persist_directory='VectorStore', embedding_function=embeddings)
    print("----------- 本地知识库加载完成 -----------")
    main_concept()
    # 创建问答系统
    
    QA_CHAIN_PROMPT = PromptTemplate.from_template("""根据下面的上下文（context）内容回答问题。
    如果你不知道答案，就回答不知道，不要试图编造答案。
    答案最多3句话，保持答案简洁。
    总是在答案结束时说”谢谢你的提问！“
    {context}
    问题：{question}
    """)
    retriever = db.as_retriever()
    qa = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        verbose=True,
        chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
    )

def main_concept():
    global db
    global mc
    QA_CHAIN_PROMPT2 = PromptTemplate.from_template("""根据下面的上下文（context）内容回答问题。
    不要试图编造答案。
    保证答案按总分结构输出。“
    {context}
    问题：{question}
    """)
    retriever2 = db.as_retriever()
    mc = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever2,
        verbose=True,
        chain_type_kwargs={"prompt": QA_CHAIN_PROMPT2}
    )
    print("----------- 摘要总结中... -----------")
    
    model_response2 = generate_response2("给出上述内容摘要，要求结果清晰")
    
    print("----------- 摘要总结完成 -----------")
    # print(model_response2)
    
    socketio.emit('mc', {'result': model_response2}, namespace='/mainconcept')
    
    autoQA()
    
    
    
def autoQA():
    global db
    global sqa
    QA_CHAIN_PROMPT3 = PromptTemplate.from_template("""根据下面的上下文（context）内容自动生成5个简答题(包括问题与答案)。
    结果格式示例如下："问题1","答案1"@"问题2","答案2"@"问题3","答案3"@"问题4","答案4"@"问题5","答案5"。
    每个简答题之间用@隔开。
    不要带有任何无关的回复。
    在每个问题与答案之间不要有多余的符号,用,隔开就行。
    {context}
    问题：{question}
    """)
    retriever3 = db.as_retriever()
    sqa = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever3,
        verbose=True,
        chain_type_kwargs={"prompt": QA_CHAIN_PROMPT3}
    )
    print("----------- 自动出题中... -----------")
    
    model_response3 = generate_response3("生成五个简答题")
    
    print("----------- 出题完成 -----------")
    print(model_response3)
    
    socketio.emit('que-ans', {'result': model_response3}, namespace='/que-ans')


#按照视频搜索所有对应数据
def get_video_data(video_name):
    try:
        # 连接到SQLite数据库
        conn = sqlite3.connect(db_file_path)
        c = conn.cursor()

        # 查询视频数据
        c.execute("SELECT * FROM videos WHERE name=?", (video_name,))
        video_data = c.fetchone()

        # 关闭连接
        conn.close()

        # 如果找到对应视频，则返回数据
        if video_data:
            return {
                "id": video_data[0],
                "name": video_data[1],
                "tags": video_data[2],
                "storage_path": video_data[3],
                "nodejson": video_data[4],
                "edgejson": video_data[5],
                "treejson": video_data[6],
                "content": video_data[7],
                "summary": video_data[8]
            }
        else:
            print("Video not found!")
            return None
    except Exception as e:
        print("Error getting video data:", e)
        return None
    
#联查：输入node返回查到的视频id(list)
def get_video_ids_by_node(node_content):
    try:
        # 连接到 SQLite 数据库
        conn = sqlite3.connect(db_file_path)
        c = conn.cursor()

        # 查询节点对应的视频 ID
        c.execute("SELECT video_id FROM link WHERE node=?", (node_content,))
        video_ids = c.fetchone()
        #输出找到的id列表
        # print(video_ids)
        # 关闭连接
        conn.close()

        # 如果找到对应视频 ID，则返回视频 ID
        if video_ids:
            return json.loads(video_ids[0])
        else:
            print("No videos found for the node:", node_content)
            return None

    except Exception as e:
        print("Error getting video IDs by node:", e)
        return None

#输入视频id（list）返回全部内容
def get_videos_data_by_ids(video_ids):
    try:
        # 连接到 SQLite 数据库
        conn = sqlite3.connect(db_file_path)
        c = conn.cursor()

        # 查询视频数据
        videos_data = []
        for video_id in video_ids:
            c.execute("SELECT * FROM videos WHERE id=?", (video_id,))
            video_data = c.fetchone()
            if video_data:
                videos_data.append({
                    "id": video_data[0],
                    "name": video_data[1],
                    "tags": video_data[2],
                    "storage_path": video_data[3],
                    "nodejson": video_data[4],
                    "edgejson": video_data[5],
                    "treejson": video_data[6],
                    "content": video_data[7]
                })

        # 关闭连接
        conn.close()

        return videos_data
    except Exception as e:
        print("Error getting videos data by IDs:", e)
        return None

def delete_files_in_folder(folder_path):
    # 遍历文件夹中的所有文件
    for file_name in os.listdir(folder_path):
        # 构建文件的绝对路径
        file_path = os.path.join(folder_path, file_name)
        try:
            # 删除文件
            if os.path.isfile(file_path):
                os.remove(file_path)
            # 如果是文件夹，则递归删除文件夹中的文件
            elif os.path.isdir(file_path):
                delete_files_in_folder(file_path)
        except Exception as e:
            print("",end='')


app=Flask(__name__)
CORS(app)

socketio = SocketIO(app, cors_allowed_origins="*")

def chat_handler(message):
    user_input = message
    print(f"用户：{user_input}")  # 打印用户输入

    model_response = generate_response(user_input)
    
    socketio.emit('reply', {'reply': model_response}, namespace='/chat')

# chat回应函数
def generate_response(user_input):
    response = qa.run(user_input)
    return json.dumps(response)

# 摘要总结函数
def generate_response2(user_input):
    response = mc.run(user_input)
    return json.dumps(response)

def generate_response3(user_input):
    response = sqa.run(user_input)
    return json.dumps(response)

@socketio.on('connect', namespace='/mainconcept')
def handle_connect():
    print('mainconcept连接已建立')
    
@socketio.on('connect', namespace='/que-ans')
def handle_connect():
    print('que-ans连接已建立')


@socketio.on('connect', namespace='/chat')
def handle_connect():
    print('socket连接已建立')

@socketio.on('disconnect', namespace='/chat')
def handle_disconnect():
    print('socket连接已关闭')

@socketio.on('message', namespace='/chat')
def handle_message(message):
    print('收到消息:', message)
    # 在这里实现处理消息的逻辑，并发送响应消息
    chat_handler(message)
    

# 与前端实时通信，确认处理进度
def process_data(data):
   
    print("----------- 当前提取进度: %d%% -----------" % data)
    socketio.emit('progress', {'progress': data}, namespace='/data')

    # 数据处理完成，发送处理完成的 JSON 数据给前端
    # data = {"result": "processed_data"}
    # socketio.emit('data_ready', data, namespace='/data')

def process_data2(data):
   
    socketio.emit('data_ready', {"result": data}, namespace='/data')
    
    
def process_data3(data):
   
    socketio.emit('data_ready2', {"result": data}, namespace='/data')

@socketio.on('connect', namespace='/data')
# def handle_connect():
#     # 当客户端连接时开始数据处理
#     socketio.start_background_task(process_data)

# 与transformer进程通信 
@socketio.on('connect', namespace='/transformer_')
def handle_connect():
    print('transformer_Client connected')

@socketio.on('transformer_message', namespace='/transformer_')
def handle_client_message(message_data):
    message = message_data['message']
    process_data(message)

@socketio.on('transformer_textpath', namespace='/transformer_')
def handle_transformer_textpath(message_data):
    message = message_data['message']
    print("----------- 本地知识库加载中... -----------")
    text_get(message)

@socketio.on('connect', namespace='/graph')
def handle_connect():
    print('graph connected')
    
@socketio.on('transformer_graph', namespace='/transformer_')
def handle_transformer_graph(message_data):
    nodes_json = message_data['nodes_json']
    edges_json = message_data['edges_json']
    # print("----------- 图谱数据生成 -----------")
       
    # print(nodes_json)
    
    socketio.emit('graph', {'nodes_json': nodes_json, 'edges_json':edges_json}, namespace='/graph')


# 与cpu_api进程通信 
@socketio.on('connect', namespace='/cpu_api')
def handle_connect():
    print('cpu_api_Client connected')

@socketio.on('cpu_api_message', namespace='/cpu_api')
def handle_client_message(message_data):
    message = message_data['message']
    process_data(message)
    
# 与generate_data进程通信 
@socketio.on('connect', namespace='/generate_data')
def handle_connect():
    print('generate_data_Client connected')

@socketio.on('generate_data_message', namespace='/generate_data')
def handle_client_message(message_data):
    message = message_data['message']
    process_data(message)

@socketio.on('generate_Jsondata_message', namespace='/generate_data')
def handle_client_message2(message_data):
    message = message_data['message']
    process_data2(message)
    
@socketio.on('generate_maincpdata_message', namespace='/generate_data')
def handle_client_message3(message_data):
    message = message_data['message']
    process_data3(message)  

    
@app.route('/init',methods=['GET','POST'])
def init():
    data=request.get_json()
    video_name=data.get('videoname')
    video_data=get_video_data(video_name)
    return_data={
        'nodejson':video_data['nodejson'],
        'edgejson':video_data['edgejson'],
        'treejson':video_data['treejson'],
        'content':video_data['content'],
        'summary':video_data['summary']
        }
    return return_data 

@app.route('/liancha',methods=['GET','POST'])
def liancha():
    data=request.get_json()
    node_name=data.get('node')
    # print(node_name)
    
    node_name = node_name.strip('"')
    
    print(node_name)
    
    node_content = "能量"
    video_data = get_videos_data_by_ids(get_video_ids_by_node(node_name))
    
    

    # video_ids=get_video_ids_by_node(node_name)
    print(video_data)
    # video_data=get_videos_data_by_ids(video_ids)
    video_data_json = json.dumps(video_data, ensure_ascii=False)

# 将 JSON 字符串写入文件中
    with open('video_data.json', 'w', encoding='utf-8') as json_file:
        json_file.write(video_data_json)
        
    video_data={'video_data':video_data}
    
    return video_data

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'video' not in request.files:
        return jsonify({'error': 'No file part'})
    
    video = request.files['video']
    
    video_name = request.form.get('videoName')
    
    video_tag = request.form.get('videoTag')
    
    # Save the file to a location
    video_path = 'temp/temp_video.mp4'
    
    audio_path = 'temp/temp_audio.mp3'

    video.save(video_path)
    
    # print("视频存储成功")
    
    
    delete_files_in_folder(r'contexts')
    
    # 调用处理视频的函数来提取音频
    transformerVideoToAudio(video_path, audio_path, video_name, video_tag)
    # subprocess.run(["python", "transformer.py", "--param1", video_path, "--param2", audio_path, "--param3", video_name, "--param4", video_tag])

    # 删除上下文contexts中的内容，便于下一个视频生成本地知识库
    delete_files_in_folder(r'contexts')
    
    
    return jsonify({'message': 'File uploaded successfully'})




# 启动 Flask 应用程序
if __name__ == '__main__':
    # app.run(debug=True,port=8090)
    socketio.run(app, debug=False, port=8090)
    
    
